import math
import time
import random
import copy
from game_rules import place_stone, mark_territory, calculate_score

# 难度配置 - 专注于得分最大化
DIFFICULTY_CONFIG = {
    "25k": [30, 2.0, 20],
    "15k": [50, 2.0, 30],
    "10k": [80, 3.0, 40],
    "5k": [120, 3.0, 50],
    "2k": [200, 5.0, 60],
    "1d": [300, 10.0, 70],
    "2d": [400, 10.0, 80],
    "3d": [500, 15.0, 90],
    "4d": [600, 15.0, 100],
    "5d": [700, 15.0, 110],
    "6d": [800, 20.0, 120],
    "7d": [900, 20.0, 130],
    "8d": [1000, 30.0, 140],
    "1p": [1200, 2.0, 150]
}


class MCTSNode:
    """MCTS树节点 - 专注于得分最大化"""

    def __init__(self, board, current_player, move=None, parent=None, komi=7.5, position_history=None):
        self.board = copy.deepcopy(board)
        self.board_size = len(board)
        self.current_player = current_player
        self.move = move
        self.parent = parent
        self.children = []
        self.total_score = 0.0  # 改为累计得分
        self.visits = 0
        self.komi = komi
        self.captures = {'BLACK': 0, 'WHITE': 0}

        # 历史局面记录
        self.position_history = copy.deepcopy(position_history) if position_history else []
        self.position_history.append(copy.deepcopy(board))

        self.untried_moves = self.get_legal_moves()

    def get_legal_moves(self):
        """获取当前节点的合法移动"""
        moves = []
        for r in range(self.board_size):
            for c in range(self.board_size):
                if self.board[r][c] is None:
                    _, _, success = place_stone(
                        self.board, self.current_player, r, c,
                        self.board_size, self.position_history
                    )
                    if success:
                        moves.append((r, c))
        moves.append(None)  # pass
        return moves

    def select_child(self):
        """使用UCT算法选择最优子节点 - 最大化得分"""

        def uct(child):
            if child.visits == 0:
                return float('inf')
            # 使用平均得分而非胜率
            return (child.total_score / child.visits) + math.sqrt(2 * math.log(self.visits) / child.visits)

        return max(self.children, key=uct)

    def add_child(self, move):
        """添加新的子节点"""
        node = MCTSNode(
            self.board,
            'WHITE' if self.current_player == 'BLACK' else 'BLACK',
            move,
            self,
            self.komi,
            self.position_history
        )

        # 执行移动
        if move is not None:
            new_board, new_captures, success = place_stone(
                self.board,
                self.current_player,
                move[0],
                move[1],
                self.board_size,
                self.position_history
            )
            if success:
                node.board = new_board
                for color in node.captures:
                    node.captures[color] = self.captures[color] + new_captures.get(color, 0)
                node.position_history.append(copy.deepcopy(new_board))

        if move in self.untried_moves:
            self.untried_moves.remove(move)
        self.children.append(node)
        return node

    def update(self, score):
        """更新节点统计信息 - 使用得分而非胜率"""
        self.visits += 1
        self.total_score += score


class MCTS:
    """蒙特卡洛树搜索AI - 专注于得分最大化"""

    def __init__(self, difficulty='5k', komi=7.5, move_count=0):
        config = DIFFICULTY_CONFIG.get(difficulty, [500, 10.0, 120])

        # 根据步数动态调整参数
        if move_count <= 20:
            # 前20手：使用上限的10%
            self.max_iterations = int(min(config[0], 5000) * 0.1)
            self.max_depth = int(config[2] * 0.1)
        elif move_count <= 40:
            # 第21-40手：使用上限的20%
            self.max_iterations = int(min(config[0], 5000) * 0.2)
            self.max_depth = int(config[2] * 0.2)
        elif move_count <= 50:
            # 第41-50手：使用上限的30%
            self.max_iterations = int(min(config[0], 5000) * 0.3)
            self.max_depth = int(config[2] * 0.3)
        elif move_count <= 80:
            # 第51-80手：使用上限的80%
            self.max_iterations = int(min(config[0], 5000) * 0.5)
            self.max_depth = int(config[2] * 0.5)
        elif move_count <= 100:
            # 第81-100手：使用上限的60%
            self.max_iterations = int(min(config[0], 5000) * 0.6)
            self.max_depth = int(config[2] * 0.6)
        elif move_count <= 150:
            # 第101-150手：使用上限的80%
            self.max_iterations = int(min(config[0], 5000) * 0.8)
            self.max_depth = int(config[2] * 0.8)
        else:
            # 150手以上：使用上限的100%
            self.max_iterations = min(config[0], 5000)
            self.max_depth = config[2]

        self.max_time = min(config[1], 240.0)
        self.komi = komi
        self.difficulty = difficulty
        self.player_color = None  # 记录当前玩家颜色
        self.move_count = move_count  # 保存当前步数
        self.surrender_threshold = -10.0  # 投降阈值（分数差）

    # 在MCTS类的get_move方法中修改进度打印部分
    def get_move(self, board, current_player, position_history):
        """获取AI的最佳移动 - 专注于最大化最终得分"""
        start_time = time.time()
        self.player_color = current_player  # 记录当前玩家颜色

        if len(board) == 0:
            return None

        # 打印当前步数和调整后的参数
        print(f"当前步数: {self.move_count}, 迭代次数: {self.max_iterations}, 模拟深度: {self.max_depth}")

        # 创建根节点
        root = MCTSNode(board, current_player, komi=self.komi, position_history=position_history)

        iteration = 0
        total_iterations = min(self.max_iterations, int(self.max_time * 100))  # 估计最大可能迭代次数

        # 计算进度更新频率（每10%或至少每10次迭代）
        update_freq = max(10, total_iterations // 10)

        while time.time() - start_time < self.max_time and iteration < self.max_iterations:
            iteration += 1
            node = root

            # 选择阶段
            while node.untried_moves == [] and node.children != []:
                node = node.select_child()

            # 展开阶段
            if node.untried_moves != []:
                move = random.choice(node.untried_moves)
                node = node.add_child(move)

            # 模拟阶段 - 返回得分而非胜率
            score = self.rollout(node)

            # 反向传播
            while node is not None:
                node.update(score)
                node = node.parent

            # 进度显示 - 增加打印频率
            if iteration % update_freq == 0 or iteration == 1:
                elapsed_time = time.time() - start_time
                progress = min(100.0, iteration / self.max_iterations * 100)
                time_left = max(0.0, self.max_time - elapsed_time)
                iters_per_sec = iteration / elapsed_time if elapsed_time > 0 else 0

                # 更详细的进度信息
                print(f"MCTS Progress: {progress:.1f}% ({iteration}/{self.max_iterations})")
                print(f"Time elapsed: {elapsed_time:.1f}s, Time left: {time_left:.1f}s")
                print(f"Speed: {iters_per_sec:.1f} iters/sec")

                # 如果已经探索了子节点，显示最佳移动信息
                if root.children:
                    best_child = max(root.children,
                                     key=lambda c: c.total_score / c.visits if c.visits > 0 else -float('inf'))
                    move_desc = "Pass" if best_child.move is None else f"({best_child.move[0]}, {best_child.move[1]})"
                    print(
                        f"Current best move: {move_desc}, Avg score: {best_child.total_score / best_child.visits:.2f}")

        # 选择最佳移动 - 基于平均得分
        if root.children:
            # 优先选择非Pass移动
            non_pass_children = [c for c in root.children if c.move is not None]
            if non_pass_children:
                best_child = max(non_pass_children,
                                 key=lambda c: c.total_score / c.visits if c.visits > 0 else -float('inf'))
            else:
                best_child = max(root.children,
                                 key=lambda c: c.total_score / c.visits if c.visits > 0 else -float('inf'))

            # 最终结果打印
            elapsed_time = time.time() - start_time
            iters_per_sec = iteration / elapsed_time if elapsed_time > 0 else 0
            move_desc = "Pass" if best_child.move is None else f"({best_child.move[0]}, {best_child.move[1]})"

            print("\n" + "=" * 50)
            print(f"MCTS completed: {iteration} iterations in {elapsed_time:.2f}s")
            print(f"Speed: {iters_per_sec:.1f} iters/sec")
            print(f"Best move: {move_desc}")
            print(f"Avg score: {best_child.total_score / best_child.visits:.2f}")
            print("=" * 50)

            # 新增：判断是否应该投降
            if self.should_surrender(best_child):
                print(f"AI认为无法取胜（平均得分：{best_child.total_score / best_child.visits:.2f}），选择投降")
                return "SURRENDER"  # 返回特殊值表示投降

            return best_child.move

        return None

    def should_surrender(self, best_child):
        """判断是否应该投降"""
        # 只有高难度AI才会考虑投降
        if self.difficulty not in ["1d", "2d", "3d", "4d", "5d", "6d", "7d", "8d", "1p"]:
            return False

        # 计算平均得分
        avg_score = best_child.total_score / best_child.visits if best_child.visits > 0 else -float('inf')

        # 如果平均得分低于投降阈值，则投降
        if avg_score < self.surrender_threshold:
            return True

        return False

    def rollout(self, node):
        """随机模拟游戏 - 返回最终得分而非胜率"""
        # 使用更精确的得分模拟
        return self.simulate_game(node, max_depth=self.max_depth)

    def simulate_game(self, node, max_depth=100):
        """增强的随机游戏模拟 - 专注于得分"""
        board = copy.deepcopy(node.board)
        current_player = node.current_player
        consecutive_pass = 0
        depth = 0
        captures = {'BLACK': 0, 'WHITE': 0}  # 初始化捕获数

        # 模拟直到终局
        while consecutive_pass < 2 and depth < max_depth:
            depth += 1
            legal_moves = self.get_legal_moves_with_score(board, current_player)

            if not legal_moves:
                consecutive_pass += 1
                current_player = 'WHITE' if current_player == 'BLACK' else 'BLACK'
                continue

            # 选择得分最高的移动
            move = max(legal_moves, key=lambda x: x[1])[0] if legal_moves else None

            if move is None:
                consecutive_pass += 1
                current_player = 'WHITE' if current_player == 'BLACK' else 'BLACK'
                continue

            # 执行移动
            new_board, new_captures, success = place_stone(
                board, current_player, move[0], move[1], len(board), []
            )

            if success:
                board = new_board
                # 更新捕获数
                for color, count in new_captures.items():
                    captures[color] += count
                consecutive_pass = 0
                current_player = 'WHITE' if current_player == 'BLACK' else 'BLACK'

        # 计算最终得分
        return self.calculate_final_score(board, node.current_player, captures)

    def get_legal_moves_with_score(self, board, player):
        """获取合法移动及其即时得分评估"""
        moves = []
        board_size = len(board)

        # 添加pass选项
        moves.append((None, self.evaluate_move(board, player, None)))

        for r in range(board_size):
            for c in range(board_size):
                if board[r][c] is None:
                    # 评估移动的即时得分
                    score = self.evaluate_move(board, player, (r, c))
                    moves.append(((r, c), score))

        return moves

    def evaluate_move(self, board, player, move):
        """评估移动的即时得分潜力"""
        if move is None:  # Pass
            return 0  # Pass通常得分为0，除非是终局

        # 创建棋盘副本
        board_copy = copy.deepcopy(board)
        r, c = move

        # 尝试落子
        new_board, captures, success = place_stone(
            board_copy, player, r, c, len(board), []
        )

        if not success:
            return -1000  # 无效移动

        # 评估得分变化
        score_before = self.estimate_score(board, player)
        score_after = self.estimate_score(new_board, player)

        # 考虑提子价值
        capture_value = sum(captures.values()) * 5

        return (score_after - score_before) + capture_value

    def estimate_score(self, board, player):
        """快速估计当前得分"""
        # 领地评估
        territory = mark_territory(board, len(board))
        black_territory = 0
        white_territory = 0

        for r in range(len(board)):
            for c in range(len(board[r])):
                if territory[r][c] == 'BLACK':
                    black_territory += 1
                elif territory[r][c] == 'WHITE':
                    white_territory += 1

        # 棋子数量评估
        black_stones = sum(row.count('BLACK') for row in board)
        white_stones = sum(row.count('WHITE') for row in board)

        # 影响力评估（基于距离中心的位置）
        center_value = 0
        center = len(board) // 2
        for r in range(len(board)):
            for c in range(len(board)):
                if board[r][c] == player:
                    # 距离中心越近价值越高
                    dist = abs(r - center) + abs(c - center)
                    center_value += max(0, 10 - dist)

        # 组合得分
        if player == 'BLACK':
            return (black_territory + black_stones * 0.5 + center_value * 0.1) + self.komi
        else:
            return (white_territory + white_stones * 0.5 + center_value * 0.1)

    def calculate_final_score(self, board, original_player, captures):
        """精确计算最终得分"""
        # 使用游戏规则中的精确计分
        territory = mark_territory(board, len(board))
        # 注意参数顺序：board, territory, komi, captures
        black_score, white_score, _ = calculate_score(board, territory, self.komi, captures)

        # 返回原始玩家的得分
        if original_player == 'BLACK':
            return black_score
        else:
            return white_score


def get_ai_move(board, board_size, current_player, difficulty, position_history, move_count):
    """根据棋盘状态获取AI移动"""
    mcts = MCTS(difficulty=difficulty, move_count=move_count)
    return mcts.get_move(board, current_player, position_history)